// Copyright (C) 2010 by Yan Huang <yhuang@virginia.edu>

package ot;

import io.NetUtils;

import java.math.*;
import java.security.SecureRandom;
import java.io.*;

import math.BitMath;

import java.util.concurrent.*;


/***************************************************************************************************
*
*
* This file is part of ME-SFE, a secure two-party computation framework.
*
* Copyright (c) 2012 - 2013 Wilko Henecka and Thomas Schneider
*
* ME-SFE is free software; you can redistribute it and/or modify it under the terms of the
* GNU General Public License as published by the Free Software Foundation; either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without
* even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* General Public License for more details.
*
* You should have received a copy of the GNU General Public License along with this program.  
* If not, see <http://www.gnu.org/licenses/>.
*
* Getting Source ==============
*
* Source for this application is maintained at code.google.com, a repository for free software
* projects.
*
* For details, please see http://code.google.com/p/me-sfe/
*/


class ThreadStep1 implements Runnable {
        BigInteger pk[], msg[];
        BigInteger[][] msgPairs;
        int numOfPairs, msgByteCount, msgBitLength, offset;
        BigInteger r,p,Cr;

        ThreadStep1(int numOfPairs, int offset, int msgByteCount, BigInteger r, BigInteger p, BigInteger Cr, BigInteger[][] msgPairs, int msgBitLength, BigInteger pk[], BigInteger msg[]) {
                this.msgByteCount = msgByteCount;
                this.numOfPairs = numOfPairs;
                this.r = r;
                this.p = p;
                this.Cr = Cr;
                this.offset = offset;
                this.msgPairs = msgPairs;
                this.msgBitLength = msgBitLength;
                this.pk = pk;
                this.msg = msg;
        }
        
        public void run() {
		BigInteger pubkey;
		ConcurrentCipher myCipher = new ConcurrentCipher();
		for (int i = 0; i < numOfPairs; i++) {
			//pk = NetUtils.readBigInteger(byteCount, dis);
			pubkey = pk[offset+i].modPow(r, p);
			msg[2*offset+2*i+0] = myCipher.encrypt(pubkey, msgPairs[offset+i][0], msgBitLength);
			//NetUtils.writeBigInteger(msg, msgByteCount, dos);
			pubkey = Cr.multiply(pubkey.modInverse(p)).mod(p);
			msg[2*offset+2*i+1] = myCipher.encrypt(pubkey, msgPairs[offset+i][1], msgBitLength);
			//NetUtils.writeBigInteger(msg, msgByteCount, dos);
		}
        }
}

public class NPOTSenderMultiThreading extends Sender {

       	private static final boolean READALL_THEN_STARTALL = true; // first read all data then start all threads

	private static SecureRandom rnd = new SecureRandom();
	private static final int certainty = 80;

	protected final static int qLength = 160;
	protected final static int pLength = 1248;

	private BigInteger p, q, g, C, r;
	private BigInteger Cr, gr;
	private int byteCount = BitMath.byteCount(pLength);
	private int msgByteCount;

	private int nrOfProcessors = Runtime.getRuntime().availableProcessors();

	public NPOTSenderMultiThreading(int numOfPairs, int msgBitLength, DataInputStream in, DataOutputStream out)
			throws Exception {
		super(numOfPairs, msgBitLength, in, out);
		msgByteCount = BitMath.byteCount(msgBitLength);
		initialize();
		//System.out.println("new NPOTSenderMultiThreading!");
	}
	
	public NPOTSenderMultiThreading(int numOfPairs, int msgBitLength, DataInputStream in, DataOutputStream out, int nOfThreads)
			throws Exception {
		super(numOfPairs, msgBitLength, in, out);
		this.nrOfProcessors = nOfThreads;
		msgByteCount = BitMath.byteCount(msgBitLength);
		initialize();
		//System.out.println("new NPOTSenderMultiThreading!");
	}

	public void execProtocol(BigInteger[][] msgPairs) throws Exception {
		super.execProtocol(msgPairs);

		step1();
	}

	private void initialize() throws Exception {
		File keyfile = new File("NPOTKey_" + pLength);
		if (keyfile.exists()) {
			try{
				readKeyFile(keyfile);
			}catch(Exception e){
				System.out.println("Error reading Keyfile. Creating a new one...");
				generateKey();
				writeKeyFile(keyfile);
			}

		} else {
			generateKey();
			writeKeyFile(keyfile);
		}
		
		NetUtils.writeBigInteger(C, dos);
		NetUtils.writeBigInteger(p, dos);
		NetUtils.writeBigInteger(q, dos);
		NetUtils.writeBigInteger(g, dos);
		NetUtils.writeBigInteger(gr, dos); 
		dos.writeInt(msgBitLength);
		dos.flush();
		
	}

	private void step1() throws Exception {
		//BigInteger[] pk0 = NetUtils.readBigIntegerArray(dis);
		//BigInteger[] pk1 = new BigInteger[numOfPairs];
		//BigInteger[][] msg = new BigInteger[numOfPairs][2];

		BigInteger pk[] = new BigInteger[numOfPairs];
		BigInteger msg[] = new BigInteger[2*numOfPairs];
		
		long start = System.currentTimeMillis();

		LinkedBlockingQueue<Runnable> queue = new LinkedBlockingQueue<Runnable>();
		ThreadPoolExecutor executor = new ThreadPoolExecutor(nrOfProcessors, nrOfProcessors, 1000, TimeUnit.MILLISECONDS, queue);

		if (READALL_THEN_STARTALL) {
                        // Read data
                        for (int i = 0; i < numOfPairs; i++) {
                                pk[i] = NetUtils.readBigInteger(byteCount, dis);
                        }
                        
                        // Start threads
                        for (int i = 0; i < nrOfProcessors; i++) {
                                executor.execute(new ThreadStep1(numOfPairs/nrOfProcessors, i*numOfPairs/nrOfProcessors, msgByteCount, r, p, Cr, msgPairs, msgBitLength, pk, msg));
                        }
                } else {
                        // Read data and start threads
                        for (int i = 0; i < nrOfProcessors; i++) {
                                for (int j=0; j<numOfPairs/nrOfProcessors; j++) {
                                        pk[i*numOfPairs/nrOfProcessors+j] = NetUtils.readBigInteger(byteCount, dis);
                                }
                                executor.execute(new ThreadStep1(numOfPairs/nrOfProcessors, i*numOfPairs/nrOfProcessors, msgByteCount, r, p, Cr, msgPairs, msgBitLength, pk, msg));
                        }
                }
		
                // collect threads
                executor.shutdown();
                try {
			executor.awaitTermination(10, TimeUnit.SECONDS);
		} catch (Exception e) {
		}

                // write data
		for (int i = 0; i < 2*numOfPairs; i++) {
                       NetUtils.writeBigInteger(msg[i], msgByteCount, dos);
		}
		dos.flush();

		//System.out.print("Time_MT: "); System.out.println(System.currentTimeMillis() - start);
	}
	
	private void readKeyFile(File keyfile) throws Exception{
		FileInputStream fin = new FileInputStream(keyfile);
		ObjectInputStream fois = new ObjectInputStream(fin);
		C = (BigInteger) fois.readObject();
		p = (BigInteger) fois.readObject();
		q = (BigInteger) fois.readObject();
		g = (BigInteger) fois.readObject();
		gr = (BigInteger) fois.readObject();
		r = (BigInteger) fois.readObject();
		Cr = (BigInteger) fois.readObject();
		fois.close();
	}
	
	private void generateKey(){
		BigInteger rr;	
		do{
			q = new BigInteger(qLength, certainty, rnd);
			rr = new BigInteger(pLength - q.bitLength(), rnd);
			p = q.multiply(rr).add(BigInteger.ONE);
		}while(!p.isProbablePrime(certainty));
		do{
			g = (new BigInteger(qLength, rnd)).mod(q);
		}while(g.modPow(BigInteger.valueOf(2), p).equals(BigInteger.ONE));
		
		g = g.modPow(BigInteger.valueOf(2), p);
		r = (new BigInteger(qLength, rnd)).mod(q);
		gr = g.modPow(r, p);
		BigInteger tmp = (new BigInteger(qLength, rnd)).mod(q);
		C = tmp.modPow(tmp, p);
		Cr = C.modPow(r, p);
	}
	
	private void writeKeyFile(File keyfile) throws Exception{
		FileOutputStream fout = new FileOutputStream(keyfile);
		ObjectOutputStream foos = new ObjectOutputStream(fout);

		foos.writeObject(C);
		foos.writeObject(p);
		foos.writeObject(q);
		foos.writeObject(g);
		foos.writeObject(gr);
		foos.writeObject(r);
		foos.writeObject(Cr);
		foos.flush();
		foos.close();
	}
}
